#!/usr/bin/env python3

##################################################################
# sjpeg converter script version 1.0
# Dependencies: (PYTHON-3)
##################################################################
SJPG_FILE_FORMAT_VERSION = "V1.00"  #
JPEG_SPLIT_HEIGHT = 16
##################################################################
import math, os, sys, time
from PIL import Image

def convert_jpg_to_sjpg(input_file, output_file_path):
    try:
        im = Image.open(input_file)
    except:
        print(f"\nFile not found or invalid image: {input_file}")
        return

    print(f"\nConverting {input_file}...\n")
    start_time = time.time()
    width, height = im.size

    print("Input:")
    print(f"\t{input_file}")
    print(f"\tRES = {width} x {height}\n")

    lenbuf = []
    block_size = JPEG_SPLIT_HEIGHT
    splits = math.ceil(height / block_size)

    sjpeg_data = bytearray()
    sjpeg = bytearray()

    row_remaining = height
    for i in range(splits):
        if row_remaining < block_size:
            crop = im.crop((0, i * block_size, width, row_remaining + i * block_size))
        else:
            crop = im.crop((0, i * block_size, width, block_size + i * block_size))

        row_remaining = row_remaining - block_size
        crop.save(str(i) + ".jpg", quality=90)

    for i in range(splits):
        with open(str(i) + ".jpg", "rb") as f:
            a = f.read()
            sjpeg_data = sjpeg_data + a
            lenbuf.append(len(a))

    header = bytearray()

    # 4 BYTES
    header = header + bytearray("_SJPG__".encode("UTF-8"))

    # 6 BYTES VERSION
    header = header + bytearray(("\x00" + SJPG_FILE_FORMAT_VERSION + "\x00").encode("UTF-8"))

    # WIDTH 2 BYTES
    header = header + width.to_bytes(2, byteorder='little')

    # HEIGHT 2 BYTES
    header = header + height.to_bytes(2, byteorder='little')

    # NUMBER OF ITEMS 2 BYTES
    header = header + splits.to_bytes(2, byteorder='little')

    # NUMBER OF ITEMS 2 BYTES
    header = header + int(JPEG_SPLIT_HEIGHT).to_bytes(2, byteorder='little')

    for item_len in lenbuf:
        # WIDTH 2 BYTES
        header = header + item_len.to_bytes(2, byteorder='little')

    sjpeg = header + sjpeg_data

    # 删除临时的 .jpg 文件
    for i in range(len(lenbuf)):
        os.remove(str(i) + ".jpg")

    # 写入 .sjpg 文件
    with open(output_file_path, "wb") as f:
        f.write(sjpeg)

    time_taken = (time.time() - start_time)

    print("Output:")
    print(f"\tTime taken = {round(time_taken, 2)} sec")
    print(f"\tbin size = {round(len(sjpeg) / 1024, 1)} KB")
    print(f"\t{output_file_path}\t(bin file)\n")

def process_folder(input_folder, output_folder):
    """
    批量处理文件夹中的所有 .jpg 文件，并将结果保存到指定的输出文件夹
    """
    if not os.path.exists(input_folder):
        print(f"Folder not found: {input_folder}")
        return

    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    jpg_files = [f for f in os.listdir(input_folder) if f.lower().endswith(".jpg")]
    if not jpg_files:
        print(f"No .jpg files found in folder: {input_folder}")
        return

    print(f"Found {len(jpg_files)} .jpg files in folder: {input_folder}\n")

    for jpg_file in jpg_files:
        input_file = os.path.join(input_folder, jpg_file)
        output_file_name = os.path.splitext(jpg_file)[0] + ".sjpg"
        output_file_path = os.path.join(output_folder, output_file_name)
        convert_jpg_to_sjpg(input_file, output_file_path)

if __name__ == "__main__":
    if len(sys.argv) == 3:
        input_folder = sys.argv[1]
        output_folder = sys.argv[2]
        process_folder(input_folder, output_folder)
    else:
        print("Usage:\n\tpython " + sys.argv[0] + " <input_folder> <output_folder>")
